package org.emailclassifier.classifiers;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;
import java.util.Vector;

import org.emailclassifier.preprocessing.DataSetDocument;

import libsvm.svm;
import libsvm.svm_model;
import libsvm.svm_node;
import libsvm.svm_parameter;
import libsvm.svm_problem;

public class SVMClassifier extends ClassifierBase {
	
	
	/**
	 * Category indexes for the probability matrix
	 */
	private HashMap<String, Integer> m_catIndex = null;
	private HashMap<Integer, String> m_catIndexi = null;
	private Vector<Double> vy = new Vector<Double>();
	private Vector<svm_node[]> vx = new Vector<svm_node[]>();
	private Map<String, Integer> termsNum = new TreeMap<String, Integer>();
	private Map<Integer,String> termsNumI = new TreeMap<Integer,String>();
	private svm_problem prob;
	private svm_model model;

	public void train() {
		
		m_catIndex = null;
		m_catIndexi = null;
		vy = new Vector<Double>();
		vx = new Vector<svm_node[]>();
		termsNum = new TreeMap<String, Integer>();
		termsNumI = new TreeMap<Integer,String>();
		
		// map of numeric term identifiers and terms
		int i = 0;
		for( String term : m_trainingSet.getGlobalVocabulary() ) {
			termsNum.put(term, i);
			termsNumI.put(i, term);
			i++;
		}
		
		// Generate indexing tables for categories
		m_catIndex = new HashMap<String, Integer>();
		m_catIndexi = new HashMap<Integer, String>();
		i = 0;
		for (String category : m_trainingSet.getCategories())
		{
			m_catIndexi.put(i, category);
			m_catIndex.put(category, i++);
		}
		
		
		// numeric category number
		int categegoryNum = 0;
		
		for( String cat : m_trainingSet.getCategories() ) {
			
			for( DataSetDocument doc : m_trainingSet.getDocumentsInCategory(cat) ) {
			
				vy.addElement( (double) categegoryNum );
				List<svm_node> svm_nodeList = new ArrayList<svm_node>();
				
				i = 0;
				for ( String term : m_trainingSet.getGlobalVocabulary() )
				{
					if (doc.getTermFrequencies().containsKey(term))
					{
						svm_node node = new svm_node();
						node.index = i;
						node.value = doc.getTermFrequencies().get(term);
						svm_nodeList.add( node );
					}
					i++;
				}
				
				/*for ( Map.Entry<String, Integer> entry : doc.getTermFrequencies().entrySet() ) {
					if ( !termsNum.containsKey(entry.getKey()) )
						continue;
					svm_node node = new svm_node();
					node.index =  termsNum.get( entry.getKey() );
					node.value = entry.getValue();
					svm_nodeList.add( node );				
				}*/
				
				svm_node x[] = svm_nodeList.toArray(new svm_node[0]);
				vx.addElement(x);

				
			}
			categegoryNum++;
		}
		
		prob = new svm_problem();
		prob.l = vy.size();
		prob.x = new svm_node[prob.l][];
		for(int k = 0; k < prob.l; k++) {
			prob.x[k] = vx.elementAt(k);
		}
		prob.y = new double[prob.l];
		for(int l=0; l<prob.l; l++) {
			prob.y[l] = vy.elementAt(l);
		}
		
		svm_parameter param = new svm_parameter();
		// default values
		param.svm_type = svm_parameter.C_SVC;
		param.kernel_type = svm_parameter.LINEAR; //.RBF;
		param.degree = 3;
		param.gamma = 1/termsNum.size();	// 1/num_features
		param.coef0 = 0;
		param.nu = 0.5;
		param.cache_size = 1000;
		param.C = 100;
		param.eps = 1e-3;
		param.p = 0.1;
		param.shrinking = 1;
		param.probability = 1;
		param.nr_weight = 0;
		param.weight_label = new int[0];
		param.weight = new double[0];
		model = svm.svm_train(prob, param);
		
		m_isTrained = true;
	}

	@Override
	public String[] determineClasses(DataSetDocument content) {

		//DataOutputStream output = null;
		//try { output = new DataOutputStream(new BufferedOutputStream(new FileOutputStream("hs.ouput"))); }
		//catch (FileNotFoundException e1) { e1.printStackTrace(); }

		List<svm_node> svm_nodeList = new ArrayList<svm_node>();
		
		int i = 0;
		for ( String term : m_trainingSet.getGlobalVocabulary() )
		{
			if (content.getTermFrequencies().containsKey(term))
			{
				svm_node node = new svm_node();
				node.index = i;
				node.value = content.getTermFrequencies().get(term);
				svm_nodeList.add( node );
			}
			i++;
		}
		
		/*	
		for ( Map.Entry<String, Integer> entry : content.getTermFrequencies().entrySet() ) {
			if ( !termsNum.containsKey(entry.getKey()) )
				continue;
			svm_node node = new svm_node();
			node.index =  termsNum.get( entry.getKey() );
			node.value = entry.getValue();
			svm_nodeList.add( node );				
		}
		*/
		svm_node x[] = svm_nodeList.toArray(new svm_node[0]);
		
		int catCount = svm.svm_get_nr_class(model);
		
		double prob_estimates[] = new double[catCount];
		
		svm.svm_predict_probability(model, x, prob_estimates);
		
		HashMap<String, Double> unsorted = new HashMap<String, Double>();
		
		for (i = 0; i<catCount; i++)
		{
			unsorted.put(m_catIndexi.get(i), prob_estimates[i]);
		}
		
		ClassProbComparator comp = new ClassProbComparator(unsorted);
		TreeMap<String, Double> sortedCategories = new TreeMap<String, Double>(comp);
		
		sortedCategories.putAll(unsorted);
		
		return sortedCategories.keySet().toArray(new String[0]);
	}
	
	public String getDescription()
	{
		return "SVM Classifier";
	}

	
	
}
